package com.ruoyi.ade.llm.chunk;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 基于标点符号的分块服务
 */
public class ChunkByMarkService implements IChunkService{
    @Override
    public List<String> chunk(String doc, IAbstractChunkRequest request) {
        ChunkByMarkRequest req = (ChunkByMarkRequest) request;
        String mark = req.getMarks();
        int windowSize = req.getWindowSize();
        int step = req.getStep();
        boolean plzh = req.isPlzh();
        int minSize = req.getMinChunckSize();

        // 1. 拆分句子并保留原始标点符号
        List<String> sentences = new ArrayList<>();
        List<String> punctuations = new ArrayList<>();
        splitSentencesWithPunctuation(doc, mark, sentences, punctuations);
        if(sentences.size()<windowSize) return Arrays.asList(new String[]{doc});

        // 2. 滑动窗口处理（使用原始标点拼接）
        List<String> windows = generateSlidingWindows(sentences, punctuations, windowSize, step);

        // 3. 排列组合处理
        List<String> result = new ArrayList<>(windows);
        if (plzh) {
            List<String> expanded = new ArrayList<>();
            for (String window : windows) {
                expanded.addAll(generateCombinations(window, mark));
            }
            result = expanded;
        }

        // 4. 去重
        Set<String> unique = new LinkedHashSet<>();
        for(String str:result){
            if(str.length()>=minSize)unique.add(str);
        }
        return new ArrayList<>(unique);

    }

    // 拆分句子并记录原始标点
    private static void splitSentencesWithPunctuation(String doc, String mark,
                                                      List<String> sentences,
                                                      List<String> punctuations) {
        String regex = "([^" + Pattern.quote(mark) + "]+)([" + Pattern.quote(mark) + "]?)";
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(doc);

        while (matcher.find()) {
            sentences.add(matcher.group(1));
            String punct = matcher.group(2);
            if (!punct.isEmpty()) {
                punctuations.add(punct);
            }
        }
    }

    /**
     *  生成滑动窗口字符串（使用原始标点）
     * @param sentences
     * @param punctuations
     * @param windowSize
     * @param step
     * @return
     */
    private static List<String> generateSlidingWindows(List<String> sentences,
                                                       List<String> punctuations,
                                                       int windowSize,
                                                       int step) {
        List<String> windows = new ArrayList<>();
        int maxStart = sentences.size() - windowSize;

        for (int start = 0; start <= maxStart; start += step) {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < windowSize; i++) {
                int pos = start + i;
                sb.append(sentences.get(pos));
                // 添加标点（最后一个句子后无标点）
                if (i < windowSize - 1 && pos < punctuations.size()) {
                    sb.append(punctuations.get(pos));
                }
            }
            windows.add(sb.toString());
        }
        return windows;
    }

    /**
     * 生成所有可能的子串组合
     * @param text
     * @param mark
     * @return
     */
    private static List<String> generateCombinations(String text, String mark) {
        String[] parts = text.split("[" + Pattern.quote(mark) + "]");
        List<String> combinations = new ArrayList<>();

        for (int i = 0; i < parts.length; i++) {
            StringBuilder sb = new StringBuilder();
            for (int j = i; j < parts.length; j++) {
                sb.append(parts[j]);
                if (j < parts.length - 1) {
                    // 从原始文本中恢复标点符号
                    String punct = findOriginalPunctuation(text, parts, j);
                    sb.append(punct);
                }
                combinations.add(sb.toString());
            }
        }
        return combinations;
    }

    /**
     * 查找原始标点符号
     * @param text
     * @param parts
     * @param index
     * @return
     */
    private static String findOriginalPunctuation(String text, String[] parts, int index) {
        int start = 0;
        for (int i = 0; i <= index; i++) {
            start += parts[i].length();
            if (i < index) {
                // 标点符号的位置在 parts[i] 末尾
                start += 1; // 标点符号占1个字符
            }
        }
        if (start < text.length()) {
            return text.substring(start, start + 1);
        }
        return "";
    }

    public static void main(String[] args) {
//        IChunkService service  = new ChunkByMarkService();
//        ChunkByMarkRequest req = new ChunkByMarkRequest();
//        req.setMarks("。，；：,?？");
//        req.setPlzh(false);
//        req.setWindowSize(3);
//        req.setStep(1);
//        System.out.println(JSONObject.toJSONString(req));
//        String text = "患者肝功能损害考虑药物性肝炎可能";
//        List<String> list = service.chunk(text,req);
//        for(String r:list){
//            System.out.println(r);
//        }
    }
}
